<!DOCTYPE html>
<html>
<head>
  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@3.18.0/dist/tf.min.js"></script>
  <!-- <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-layers@3.18.0/dist/tf-layers.min.js"></script>
  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-datasets@2.8.1/dist/tf-datasets.min.js"></script> -->
</head>
<body>
  <button id="trainBtn">Train Model</button>
  <div id="results"></div>

  <script>
    document.getElementById('trainBtn').addEventListener('click', async () => {
      const { train, test } = await MNIST.load();
      
      // 数据预处理
      const xs = tf.tidy(() => {
        return train.xs.map(x => tf.tensor2d(x, [28, 28]).div(255));
      });
      const ys = train.labels;

      // 构建模型
      const model = tf.sequential({
        layers: [
          tf.layers.conv2d({
            inputShape: [28, 28, 1],
            filters: 32,
            kernelSize: 5,
            strides: 1,
            activation: 'relu'
          }),
          tf.layers.maxPooling2d({
            poolSize: [2, 2],
            strides: [2, 2]
          }),
          tf.layers.conv2d({
            filters: 64,
            kernelSize: 5,
            strides: 1,
            activation: 'relu'
          }),
          tf.layers.maxPooling2d({
            poolSize: [2, 2],
            strides: [2, 2]
          }),
          tf.layers.flatten(),
          tf.layers.dense({units: 10, activation: 'softmax'})
        ]
      });

      // 训练模型
      const optimizer = tf.train.adam(0.001);
      model.compile({optimizer, loss: 'categoricalCrossentropy'});

      for (let i = 0; i < 10; i++) {
        await model.fit(xs, ys, {epochs: 1, batchSize: 128});
        console.log(`Epoch ${i + 1} completed`);
      }

      // 评估模型
      const testXs = tf.tidy(() => {
        return test.xs.map(x => tf.tensor2d(x, [28, 28]).div(255));
      });
      const testYs = test.labels;

      const lossAndAccuracy = await model.evaluate(testXs, testYs);
      console.log('Test Loss:', lossAndAccuracy.dataSync());
      console.log('Test Accuracy:', lossAndAccuracy.dataSync()‌:ml-citation{ref="1" data="citationList"});

      document.getElementById('results').innerHTML = `
        <h3>Training Results</h3>
        <p>Test Loss: ${lossAndAccuracy.dataSync().toFixed(4)}</p>
        <p>Test Accuracy: ${lossAndAccuracy.dataSync()‌:ml-citation{ref="1" data="citationList"}.toFixed(4)}</p>
      `;
    });
  </script>
</body>
</html>
